import httpx
import json
import logging
from PyCmpltrtok.common import md5, sep
from python_basic.exercise.lib.log.learn_from_fastchat.logger_util import build_logger

logger = build_logger('fastapi_interface', 'fastapi_interface.log', logger_level=logging.INFO)


def text_gen(xinput, username='', dport=7750, timeout=99999999):
    xjson = dict()
    xjson['xinput'] = xinput
    xjson['username'] = username
    # xjson = dict(query=xjson)

    base_url = f'http://127.0.0.1:{dport}'
    print(base_url)
    url = '/stream_chat'
    print(url)
    entire_url = f'{base_url}{url}'
    print(entire_url)
    
    # res = requests.post(url, json=xjson, stream=True)
    # for chunk in res.iter_lines(decode_unicode=False, delimiter=b'\0'):
    #     if not chunk:
    #         continue
    #     chunk = chunk.decode('utf8')
    #     data = json.loads(chunk)
    #     xouput = data['output']
    #     yield xouput
    
    _client = httpx.Client(base_url=base_url, timeout=timeout)
    response = _client.stream("POST", url, json=xjson, data=None)
    as_json = True
    
    try:
        with response as r:
            for chunk in r.iter_text(None):
                if not chunk: # fastchat api yield empty bytes on start and end
                    continue
                if as_json:
                    try:
                        if chunk.startswith("data: "):
                            data = json.loads(chunk[6:-2])
                        elif chunk.startswith(":"): # skip sse comment line
                            continue
                        else:
                            data = json.loads(chunk)
                        yield data['output']
                    except Exception as e:
                        msg = f"接口返回json错误： ‘{chunk}’。错误信息是：{e}。"
                        logger.error(f'{e.__class__.__name__}: {msg}', exc_info=e)
                else:
                    # print(chunk, end="", flush=True)
                    yield chunk
    except httpx.ConnectError as e:
        msg = f"无法连接API服务器，请确认 ‘api.py’ 已正常启动。({e})"
        logger.error(msg)
        yield {"code": 500, "msg": msg}
    except httpx.ReadTimeout as e:
        msg = f"API通信超时，请确认已启动FastChat与API服务（详见Wiki '5. 启动 API 服务或 Web UI'）。（{e}）"
        logger.error(msg)
        yield {"code": 500, "msg": msg}
    except Exception as e:
        msg = f"API通信遇到错误：{e}"
        logger.error(f'{e.__class__.__name__}: {msg}', exc_info=e)
        yield {"code": 500, "msg": msg}


if '__main__' == __name__:
    user = 'u0001'
    gen = text_gen('怎么滑雪？', username=user)
    for x in gen:
        print(x)
        